/**
 * @file       rw_lock.hpp
 * @author     WMQ
 * @date       2023/9/9
 * @brief      atomic实现读写锁，写优先
 */

#ifndef RW_LOCK_HPP
#define RW_LOCK_HPP

#include <atomic>
#include <thread>

/**
 * @brief 读写锁类 RWLock
 */
class RWLock
{
public:
    /**
     * @brief 获取读锁，读锁可重入
     */
    void ReadLock()
    {
        uint32_t retry_times = 0;
        for (;;)
        {
            int64_t readers = lock_.load(std::memory_order_acquire);
            if (!is_write_ && readers != RWLOCK_WLOCK &&
                lock_.compare_exchange_strong(readers, readers + 1, std::memory_order_release,
                                              std::memory_order_relaxed))
            {
                return;
            }

            if (++retry_times == MAX_RETRY_TIMES)
            {
                std::this_thread::yield();
                retry_times = 0;
            }
        }
    }

    /**
     * @brief 获取写锁，写优先
     */
    void WriteLock()
    {
        uint32_t retry_times = 0;
        is_write_ = true;
        for (;;)
        {
            int64_t lock = lock_.load(std::memory_order_acquire);
            if (0 == lock &&
                lock_.compare_exchange_strong(lock, RWLOCK_WLOCK, std::memory_order_release,
                                              std::memory_order_relaxed))
            {
                is_write_ = false;
                return;
            }

            if (++retry_times == MAX_RETRY_TIMES)
            {
                std::this_thread::yield();
                retry_times = 0;
            }
        }
    }

    /**
     * @brief 释放读锁
     */
    void ReadUnlock()
    {
        lock_.fetch_sub(1, std::memory_order_release);
    }

    /**
     * @brief 释放写锁
     */
    void WriteUnlock()
    {
        lock_.fetch_add(1, std::memory_order_release);
    }

    RWLock() : lock_(0) {}

    RWLock(const RWLock &) = delete;
    RWLock &operator=(const RWLock &) = delete;

private:
    static const int64_t RWLOCK_WLOCK = -1; /**< 写锁状态 */
    static const uint32_t MAX_RETRY_TIMES = 3; /**< 让出cpu时间片之前尝试获取读写锁的最大次数 */
    std::atomic<int64_t> lock_;                /**< 读写锁: -1:锁 0:空闲 >0:读锁 */
    bool is_write_ = false;                    /**< 是否尝试获取写锁，用于写优先 */
};

#endif